package de.unibi.comet.hmm;

import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;

import de.unibi.comet.fa.Partition;
import de.unibi.comet.util.LogSpace;
import de.unibi.comet.util.Pair;

public class HMM {
	private List<? extends Node> nodes;
	private boolean logarithmic;
	
	public HMM(List<? extends Node> nodes) {
		this.nodes=nodes;
		this.logarithmic=false;
	}

	/** Calculates the new joined distribution of state and accumulated mass
	 *  after proceeding one step (one transition). */
	public StateAndMassDistribution step(StateAndMassDistribution dist) {
		if (dist.isLogarithmic()!=logarithmic) throw new IllegalArgumentException();
		if (dist.getStateCount()!=nodes.size()) {
			throw new IllegalArgumentException("Invalid number of states");
		}
		int massBound = dist.getMassBound();
		StateAndMassDistribution new_dist = dist.emptyClone();
		for (int state=0; state<nodes.size(); ++state) {
			Node node = nodes.get(state);
			// get iterator over accumulated masses in current state
			Iterator<Emission> iter = dist.getMassDistribution(state);
			while (iter.hasNext()){
				Emission so_far_acc_mass = iter.next();
				for (Transition t : node.getTransitions()) {
					// iterate over emission distribution of target node
					for (Emission e : nodes.get(t.getTargetIdx()).getEmissionDistribution()) {
						// calculate probability addition:
						// so_far_acc_mass.getProbability(): prob of being in state "state"
						//    and having accumulated mass "so_far_acc_mass.getMass()"
						// t.getProbability(): probability of transition from state "state"
						//    into state "t.getTargetIdx()"
						// e.getProbability(): probability for state "t.getTargetIdx()" to
						//    emit mass "e.getMass()"
						if (!logarithmic) {
							double p = so_far_acc_mass.getProbability()*t.getProbability()*e.getProbability();
							int m = so_far_acc_mass.getMass()+e.getMass();
							if ((massBound>=0) && (m>massBound)) m=massBound;
							new_dist.add(t.getTargetIdx(), m, p);
						} else {
							double p = so_far_acc_mass.getProbability()+t.getProbability()+e.getProbability();
							int m = so_far_acc_mass.getMass()+e.getMass();
							if ((massBound>=0) && (m>massBound)) m=massBound;
							new_dist.add(t.getTargetIdx(), m, p);
						}
					}
				}
			}
		}
		return new_dist;
	}
	
	/** Performs as many steps (like in step()) until convergence of state distribution is reached.
	 *  @return Returns a pair of 1) new distribution and 2) number of steps done.
	 */
	public Pair<StateAndMassDistribution, Integer> stepToEquilibrium(StateAndMassDistribution dist, int maxSteps, double accuracy) {
		if (logarithmic) accuracy=Math.log(accuracy);
		int stepsDone = 0;
		int checkInterval = 15;
		int n = 2;
		double[] stateDist1 = null;
		double[] stateDist2 = null;
		while (stepsDone<maxSteps) {
			dist = step(dist);
			++stepsDone;
			// check if convergence was reached
			if (n%checkInterval==0) {
				stateDist1=dist.getStateDistribution();
			}
			if (n%checkInterval==1) {
				stateDist2=dist.getStateDistribution();
				double maxDiff = logarithmic?Double.NEGATIVE_INFINITY:0.0;
				for (int j=0; j<stateDist1.length; ++j) {
					double diff;
					if (!logarithmic) {
						diff = Math.abs(stateDist1[j]-stateDist2[j]);
					} else {
						if (stateDist1[j]>stateDist2[j]) {
							diff = LogSpace.logSub(stateDist1[j],stateDist2[j]);
						} else {
							diff = LogSpace.logSub(stateDist2[j],stateDist1[j]);
						}
					}
					if (diff>maxDiff) maxDiff=diff;
				}
				System.out.println(String.format("i=%d, maxDiff=%g", stepsDone, maxDiff));
				if (maxDiff<accuracy) break;
			}
			++n;
		}
		return new Pair<StateAndMassDistribution, Integer>(dist, stepsDone);
	}
	
	
	public double[] convolution(StateAndMassDistribution dist, int steps) {
		if (dist.isLogarithmic()!=logarithmic) throw new IllegalArgumentException();
		if (dist.getStateCount()!=nodes.size()) {
			throw new IllegalArgumentException("Invalid number of states");
		}
		int massBound = dist.getMassBound();
		if (massBound<0) throw new IllegalArgumentException("Convolution only implemented for bounded masses");
		double[] massDistribution = new double[massBound+1];
		for (int mass=0; (mass<=dist.getMaxMass()) && (mass<=massBound); ++mass) {
			massDistribution[mass]=dist.getMassProbability(mass);
		}
		double[] stateDist = dist.getStateDistribution();
		// get emission distribution in equilibrium state
		double[] emission = getEmissionDistribution(stateDist);
		// eventually merge table entries
		if (emission.length>massBound+1) {
			double[] a = new double[massBound+1];
			System.arraycopy(emission, 0, a, 0, massBound);
			for (int j=massBound; j<emission.length; ++j) {
				if (!logarithmic) {
					a[massBound]+=emission[j];
				} else {
					a[massBound]=LogSpace.logAdd(a[massBound], emission[j]);
				}
			}
			emission=a;
		}
		// prepare table of additive emissions
		// additiveEmission[i] is the probability of emitting i or more matches
		int m = emission.length;
		double[] additiveEmission = new double[m];
		if (logarithmic) {
			for (int i=0; i<additiveEmission.length; ++i) additiveEmission[i]=Double.NEGATIVE_INFINITY;
		}
		additiveEmission[m-1]=emission[m-1];
		for (int i=m-2; i>=0; --i) {
			if (!logarithmic) {
				additiveEmission[i]=emission[i]+additiveEmission[i+1];
			} else {
				additiveEmission[i]=LogSpace.logAdd(emission[i],additiveEmission[i+1]);
			}
		}
		// main loop
		for (int stepsDone=0; stepsDone<steps; ++stepsDone) {
			double[] newMassDistribution = new double[massBound+1];
			if (logarithmic) {
				for (int i=0; i<newMassDistribution.length; ++i) newMassDistribution[i]=Double.NEGATIVE_INFINITY;
			}
			for (int i=0; i<massDistribution.length; ++i) {
				for (int j=0; j<emission.length; ++j) {
					if (i+j==massBound) {
						if (!logarithmic) {
							newMassDistribution[i+j]+=massDistribution[i]*additiveEmission[j];
						} else {
							newMassDistribution[i+j]=LogSpace.logAdd(newMassDistribution[i+j], massDistribution[i]+additiveEmission[j]);
						}
						break;
					}
					if (!logarithmic) {
						newMassDistribution[i+j]+=massDistribution[i]*emission[j];
					} else {
						newMassDistribution[i+j]=LogSpace.logAdd(newMassDistribution[i+j], massDistribution[i]+emission[j]);
					}
				}
			}
			massDistribution=newMassDistribution;
		}
		return massDistribution;
	}

	/** Returns the distribution of emissions for a given state distribution. 
	 *  NOTE: if HMM is logarithmic, a logarithmic state distribution has to be passed.*/
	private double[] getEmissionDistribution(double[] stateDistribution) {
		if (stateDistribution.length!=nodes.size()) throw new IllegalArgumentException();
		ArrayList<Double> dist = new ArrayList<Double>();
		int i=0;
		for (Node n : nodes) {
			double p = stateDistribution[i++];
			for (Emission e : n.getEmissionDistribution()) {
				int m = e.getMass();
				while (m>=dist.size()) {
					if (!logarithmic) {
						dist.add(0.0);
					} else {
						dist.add(Double.NEGATIVE_INFINITY);
					}
				}
				if (!logarithmic) {
					dist.set(m, dist.get(m)+p*e.getProbability());
				} else {
					dist.set(m, LogSpace.logAdd(dist.get(m),p+e.getProbability()));
				}
			}
		}
		double[] result = new double[dist.size()];
		i=0;
		for (double d : dist) result[i++]=d;
		return result;
	}
	
	/** Converts all probabilities to the logarithmic domain. */
	public void convertToLogSpace() {
		if (logarithmic) throw new IllegalStateException("Already logarithmic!");
		for (Node node : nodes) {
			for (Transition t : node.getTransitions()) t.convertToLogSpace();
			for (Emission e : node.getEmissionDistribution()) e.convertToLogSpace();
		}
		logarithmic=true;
	}

	/** Joines all states that have the same emission. The transition probabilities are modeled
	 *  according to given distribution (which should be the stationary distribution).
	 *  @return Returns the distribution translated into the new HMM. */
	public StateAndMassDistribution collapse(StateAndMassDistribution dist) {
		// TODO logarithmic!
		if (logarithmic) throw new IllegalStateException("Not yet implemented for logarithmic case!");
		// for each state the list of its emission
		List<Integer> emissions = new ArrayList<Integer>();
		for (Node n : nodes) {
			if (n.getEmissionDistribution().size()>1) throw new IllegalStateException("Only deterministic hmms can be collapsed!");
			int e = 0;
			if (n.getEmissionDistribution().size()==1) e=n.getEmissionDistribution().get(0).getMass();
			emissions.add(e);
		}
		// partition the state space according to the emissions 
		Partition partition = new Partition(emissions);
		// create new distribution
		int maxMass = dist.getMaxMass();
		int newStateCount = partition.blockCount();
		ArrayStateAndMassDistribution newDist = new ArrayStateAndMassDistribution(newStateCount, maxMass, dist.getMassBound(), logarithmic);
		for (int mass=0; mass<=maxMass; ++mass) {
			for (int newState=0; newState<newStateCount; ++newState) {
				double p = 0.0;
				for (int oldState : partition.block(newState)) {
					p+=dist.get(oldState, mass);
				}
				newDist.set(newState, mass, p);
			}
		}
		// create new nodes
		List<Node> newNodes = new ArrayList<Node>(partition.blockCount());
		// for every set in the parition, create one state for new hmm
		for (int i=0; i<partition.blockCount(); ++i) {
			double[] nodeTransitions = new double[partition.blockCount()];
			int e = -1;
			for (int sourceIdx : partition.block(i)) {
				if (e==-1) e=emissions.get(sourceIdx);
				double sourceProb = dist.getStateProbability(sourceIdx);
				for (Transition t : nodes.get(sourceIdx).getTransitions()) {
					nodeTransitions[partition.getBlockIndex(t.getTargetIdx())]+=t.getProbability()*sourceProb;
				}
			}
			List<Transition> transitionList = new ArrayList<Transition>();
			double newStateProb = newDist.getStateProbability(i);
			for (int j=0; j<nodeTransitions.length; ++j) {
				if (nodeTransitions[j]>0.0) transitionList.add(new Transition(j,nodeTransitions[j]/newStateProb));
			}
			List<Emission> emissionList = new ArrayList<Emission>(1);
			emissionList.add(new Emission(e,1.0));
			newNodes.add(new StandardNode(transitionList,emissionList));
		}
		nodes=newNodes;
		return newDist;
	}
	
	/** Outputs a textual representation. */
	@Override
	public String toString() {
		// create transition matrix
		double[][] transitionMatrix = new double[nodes.size()][nodes.size()];
		int k=0;
		for (Node node : nodes) {
			for (Transition t : node.getTransitions()) {
				transitionMatrix[k][t.getTargetIdx()]+=t.getProbability();
			}
			++k;
		}
		// print to string buffer
		StringBuffer sb = new StringBuffer();
		sb.append(String.format("HMM: states: %d, transition matrix:\n", nodes.size()));
	
		k=0;
		for (Node node : nodes) {
			sb.append(String.format(" %4d |", k));
			for (Transition t : node.getTransitions()) {
				sb.append(String.format(" %d:(%e) |", t.getTargetIdx(), t.getProbability()));
			}
			++k;
			sb.append("\n");
		}
//		for (int i=0; i<nodes.size(); ++i) {
//			sb.append(String.format(" %4d |", i));
//			for (int j=0; j<nodes.size(); ++j) {
//				sb.append(String.format(" %8f ", transitionMatrix[i][j]));
//			}
//			sb.append("\n");
//		}
		
		sb.append("output distribution:\n");
		k=0;
		for (Node node : nodes) {
			sb.append(String.format(" %4d |", k));
			for (Emission e : node.getEmissionDistribution()) {
				sb.append(String.format(" %d:(%8f) |", e.getMass(), e.getProbability()));
			}
			++k;
			sb.append("\n");
		}
		
		return sb.toString();
	}
	
	public int getStateCount() {
		return nodes.size();
	}
}
